'''
Author: zhuyuejiang
Date: 2021-03-16 14:09:16
LastEditTime: 2021-05-19 02:14:46
Description: ReID算法处理接口类，用于管理所有ReID算法包启动、运行等流程
'''
import os
import cv2
import time
import json
import os.path as osp
import importlib
import sys
import FrameWork.utils as reid_utils
# from .default_config import default_config
from yacs.config import CfgNode as CN

'''
@desc: 以默认参数设置完成所有配置的初始化
'''
def default_config():
    cfg = CN()
    #algorithms
    cfg.name = 'algorithms'

    print("default configing--------")
    # model
    cfg.model = CN()
    cfg.model.name = 'baseline'
    cfg.model.device = 'cuda'
    cfg.model.gpu_devices ='0'
    cfg.model.pretrained = True # automatically load pretrained model weights if available
    cfg.model.load_weights = ''# path to model weights
    cfg.model.resume = '' #path to checkpoint for resume training
    cfg.model.freeze_layers = []

    # backbone
    cfg.model.base = CN()
    cfg.model.base.name = 'resnet50'
    cfg.model.base.last_stride = 1

    # If use ImageNet pretrain model
    cfg.model.base.pretrain = False

    # loss
    cfg.model.loss = CN()
    cfg.model.loss.name = ("CrossEntropyLoss",)
    # Cross Entropy Loss options
    cfg.model.loss.ce = CN()
    cfg.model.loss.ce.epsilon = 0.1 # use label smoothing regularizer, 0 means no label smooth
    cfg.model.loss.ce.weight = 1.0
    # Triplet Loss options
    cfg.model.loss.triplet = CN()
    cfg.model.loss.triplet.margin = 0.3 # distance margin
    cfg.model.loss.triplet.weight = 1. # weight to balance hard triplet loss
    cfg.model.loss.triplet.hard_mining = False

    # data
    cfg.data = CN()
    cfg.data.root = 'reid-data'
    cfg.data.sources = 'market1501'
    cfg.data.targets = 'market1501'
    cfg.data.workers = 4 # number of data loading workers
    cfg.data.split_id = 0 # split index
    cfg.data.height = 256 # image height
    cfg.data.width = 128 # image width
    cfg.data.combineall = False # combine train, query and gallery for training
    cfg.data.transforms = ['random_flip', 'random_erase', 'random_crop'] # data augmentation
    cfg.data.padding = 0 # transform padding
    # random erasing
    cfg.data.rea = CN()
    cfg.data.rea.prob = 0.5
    cfg.data.rea.mean = [0.596, 0.558, 0.497]
    # Random Patch
    cfg.data.rpt = CN()
    cfg.data.rpt.prob = 0.5

    cfg.data.norm_mean = [0.485, 0.456, 0.406] # default is imagenet mean
    cfg.data.norm_std = [0.229, 0.224, 0.225] # default is imagenet std
    cfg.data.save_dir = 'log' # path to save log
    
    # sampler
    cfg.sampler = CN()
    cfg.sampler.name = '' # random_identity;sampler
    cfg.sampler.num_instances = 4 # number of instances per identity for RandomIdentitySampler

    # train
    cfg.train = CN()
    cfg.train.optim = 'adam'
    cfg.train.lr = 0.0003
    cfg.train.weight_decay = 5e-4
    cfg.train.weight_decay_bias = 5e-4
    cfg.train.max_epoch = 60
    cfg.train.start_epoch = 0
    cfg.train.batch_size = 32
    cfg.train.staged_lr = False # set different lr to different modules
    cfg.train.new_layers = ['classifier'] # newly added modules with default lr
    cfg.train.base_lr_mult = 0.1 # learning rate multiplier for base modules
    cfg.train.lr_scheduler = 'single_step'
    cfg.train.stepsize = [20] # stepsize to decay learning rate
    cfg.train.warmup_factor = 0.1
    cfg.train.warmup_iters = 10
    cfg.train.freeze_iters = 0
    cfg.train.warmup_method = 'linear'
    cfg.train.gamma = 0.1 # learning rate decay multiplier
    cfg.train.delay_iters = 30
    cfg.train.eta_min_lr = 0.
    cfg.train.print_freq = 50 # print frequency
    cfg.train.seed = 1 # random seed
    cfg.train.opti_open_init = False
    cfg.train.metric = "euclidean" #['euclidean','k-reciprocal','cosine']
    cfg.train.lambda_value = 0.1

    # optimizer
    cfg.train.sgd = CN()
    cfg.train.sgd.momentum = 0.9 # momentum factor
    cfg.train.sgd.dampening = 0. # dampening for momentum
    cfg.train.sgd.nesterov = False # Nesterov momentum
    cfg.train.rmsprop = CN()
    cfg.train.rmsprop.alpha = 0.99 # smoothing constant
    cfg.train.adam = CN()
    cfg.train.adam.beta1 = 0.9 # exponential decay rate for first moment
    cfg.train.adam.beta2 = 0.999 # exponential decay rate for second moment

    # test
    cfg.test = CN()
    cfg.test.batch_size = 100
    cfg.test.dist_metric = 'euclidean' # distance metric, ['euclidean', 'cosine']
    cfg.test.print_freq = 200
    cfg.test.normalize_feature = False # normalize feature vectors before computing distance
    cfg.test.ranks = [1, 5, 10, 20] # cmc ranks
    cfg.test.evaluate = False # test only
    cfg.test.eval_freq = -1 # evaluation frequency (-1 means to only test after training)
    cfg.test.start_eval = 0 # start to evaluate after a specific epoch
    cfg.test.start_save = 0
    cfg.test.rerank = False # use person re-ranking
    cfg.test.lambda_value = 0.3
    cfg.test.visrank = False # visualize ranked results (only available when cfg.test.evaluate=True)
    cfg.test.visrank_topk = 10 # top-k ranks to visualize
    cfg.test.visactmap = False # visualize CNN activation maps
    cfg.test.flip = False #test when sample flip
    cfg.test.use_cython = False
    return cfg


class ReidProcessor:
    def __init__(self, result_path):
        self._result_path = result_path
    
    
    '''
    @description: 算法配置参数初始化
    @param {*} config_file 用户指定的配置文件路径
    @param {*} model_path 模型文件路径
    @param {*} data_path 数据集路径
    @param {*} save_dir 输出结果保存目录
    @return {*}
    '''
    def init_config(self, config_file, model_path, data_path, save_dir):
        #如果配置文件已指定且存在，则从配置文件读取配置参数，否则采用默认值
        cfg = default_config()
        if config_file is not "":
            cfg.merge_from_file(config_file)

        if model_path is not "":
            cfg.model.load_weights = model_path

        if data_path is not "":
            p,f = osp.split(data_path)
            cfg.data.root = p
        print("[Info:] data root: %s" % cfg.data.root)

        if save_dir == '':
            cfg.data.save_dir = self._result_path
        else:
            cfg.data.save_dir = save_dir

        #gpu相关参数设置
        if cfg.model.device == 'cuda':
            os.environ['CUDA_VISIBLE_DEVICES'] = cfg.model.gpu_devices
            print("[Info:] Currently using GPU {}".format(os.environ['CUDA_VISIBLE_DEVICES']))
        
        return cfg

    
    '''
    @description: 在数据集上运行ReID算法
    @param {*} library_path 算法库根目录
    @param {*} model_path 模型文件路径
    @param {*} config_path 配置参数文件路径
    @param {*} input_path 数据集路径
    @param {*} query_image 用于比对检索的图片路径
    @param {*} output_path 输出结果保存目录
    @return {*}
    '''
    def run_eval(self, library_path,model_path,config_path,input_path, query_image, output_path):
        cfg = self.init_config(config_file = config_path, model_path = model_path, data_path = input_path, save_dir=output_path)

        # import lib
        sys.path.append(library_path)
        if sys.modules.get('PersonReID') is not None:
            del sys.modules['PersonReID']
        if sys.modules.get('PersonReID.reidentification') is not None:
            del sys.modules['PersonReID.reidentification']
        reid_lib = importlib.import_module('PersonReID.reidentification')

        #mkdir 
        reid_utils.mkdir_if_missing(output_path)

        #运行算法
        reid = reid_lib.ZJLabReID(cfg)
        query_img, gallery_imgs, rank, mAP = reid.run(query_image, 30)

        #保存结果
        total_gallery = len(gallery_imgs)
        current = 0
        if total_gallery > 0:
            for i in range(total_gallery):
                cv2.imwrite(osp.join(cfg.data.save_dir, 'gallery'+str(i)+'.jpg'), gallery_imgs[i])
                current += 1
                print('PlatformProcessNum:', int(current * 100 / total_gallery))

        sys.path.remove(library_path)
        print('PlatformSummary:', '{"rank":%s, "mAP": %s}' % (rank, mAP))
    

if __name__ == '__main__':
    print("[Info:] reidrunner recieved task, ReID process starting............................................")
    print("[Info:] --------------query image file:", sys.argv[1])
    print("[Info:] --------------reid library path:", sys.argv[2])
    print("[Info:] --------------reid model file:", sys.argv[3])
    print("[Info:] --------------config file:", sys.argv[4])
    print("[Info:] --------------dataset path:", sys.argv[5])
    print("[Info:] --------------result root dir:", sys.argv[6])
    print("[Info:] --------------data type:", sys.argv[7])
    store_path = "/data/result/reid"
    reid = ReidProcessor(store_path)
    if int(sys.argv[7]) == 0 :
        print("[Err:] video is unsupported by reid algorithms")
    else:
        if len(sys.argv[6]) > 0 :
            store_path = sys.argv[6]
        reid.run_eval(sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5], sys.argv[1], store_path)